require 'nn'
require 'image'

class_id = {}
i = 0
local function lines_from(file)
-- get all lines from file
  if not paths.filep(file) then return {} end
  local lines = {}
  for line in io.lines(file) do 
    label = 
    table.insert(lines,line)
    i = i + 1
    class_id[string.sub(line,1,9)] = i
  end
  return lines
end

classes_name_imagenet = lines_from('synset.txt')

maps = torch.load('cams_topk.t7')
preds = torch.load('pred_topk.t7')
clsid = torch.load('./gen/imagenet.t7')

clsid = clsid['classList']

fd = io.open('predBox_topk.txt', 'w')
count = 0
imsize = 224
debug = false
for key, value in pairs(maps) do
    count = count + 1
    folder, image_name = key:match("([^,]+)/([^,]+)")
    
    img = image.load(string.format('/root/Dataset/ILSVRC2012/val_caffe/%s',image_name))
    str = string.format('%d',tonumber(string.sub(image_name,16,23)))
    
    for k=1,5 do      
        pr = preds[key]
        class_words = clsid[pr[{1,k}]+1]     
        va = value[{k,{}}]:squeeze()
        
        map = image.scale(va:float(),(#img)[3],(#img)[2],'bicubic')
        map = (map - map:min()) / (map:max() - map:min())
        th = map:mean()
--         th = (map:mean() + map:max())/2
        a = map:ge(th)
        x = a:nonzero()
        r = x[{{},1}]
        c = x[{{},2}]
--         if debug then
--             print(image_name)
--             print(#img)
--     --         img = image.scale(img,imsize,imsize,'bicubic')
--     --         itorch.image(img)

--             cam = getCAM(img, value:float(), (#img)[3], (#img)[2])
--             itorch.image(cam)
--     --         print(c:min(), r:min(), c:max(), r:max())
--             im_box = image.drawRect(cam, c:min(), r:min(), c:max(), r:max(), {lineWidth = 1, color = {0, 255, 0}})
--             itorch.image(im_box)

--         end
    x1 = c:min() --or 1
    y1 = r:min() --or 1
    x2 = c:max() --or (#img)[3]
    y2 = r:max() --or (#img)[2]
    -- print(key)
    -- print(pr)
    -- print(class_id[class_words])

    res = string.format(' %d %d %d %d %d', class_id[class_words], x1, y1, x2, y2)
    str = str..res
    end
    fd:write(string.format('%s\n', str))
    -- if count == 100 then
    --     break
    -- end
end
fd:close()